[Minor] Add speech models#232
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 30616d7889
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| whisper_type_to_module_mapping = { | ||
| "block_input": ("encoder.layers[%s]", CONST_INPUT_HOOK), | ||
| "block_output": ("encoder.layers[%s]", CONST_OUTPUT_HOOK), | ||
| "mlp_activation": ("encoder.layers[%s].activation_fn", CONST_OUTPUT_HOOK), |
There was a problem hiding this comment.
In HuggingFace Whisper, encoder.layers[*].activation_fn is a plain callable from ACT2FN, not an nn.Module. get_module_hook later calls register_forward_hook on the resolved object (see pyvene/models/modeling_utils.py:get_module_hook), so registering mlp_activation will raise an AttributeError at runtime whenever a user selects that intervention. Consider hooking a real module (e.g., fc1 output) or wrapping the activation in an nn.Module before registering.
Useful? React with 👍 / 👎.
| "ffn1_activation": ("encoder.layers[%s].ffn1.intermediate_act_fn", CONST_OUTPUT_HOOK), | ||
| "ffn1_output": ("encoder.layers[%s].ffn1", CONST_OUTPUT_HOOK), | ||
| "ffn1_input": ("encoder.layers[%s].ffn1", CONST_INPUT_HOOK), | ||
| "ffn2_activation": ("encoder.layers[%s].ffn2.intermediate_act_fn", CONST_OUTPUT_HOOK), |
There was a problem hiding this comment.
ffn1.intermediate_act_fn and ffn2.intermediate_act_fn are callables rather than nn.Modules in the HF implementation, so get_module_hook (which calls register_forward_hook) will fail when users request ffn*_activation interventions. This makes those intervention points unusable and will raise at hook registration time.
Useful? React with 👍 / 👎.
|
@taka-yamakoshi I think codex is wrong since we hook onto act_fn elsewhere too with no problems. I'm going to merge, but do double check that that works just in case I am wrong! |
|
I checked it just now, and it seems to work just fine! |
2 participants
Description
Added speech models (Whisper & Wav2Vec-BERT).
Allow batched multi-token interventions.
Testing Done
Passed unit tests.
Checklist:
[Your Priority] Your Title